package main

import (
	"flag"
	"fmt"
	"sync"
	"time"

	"github.com/nsqio/go-nsq"
)

var (
	resultChan = make(chan int64, 10000)
	wg         = sync.WaitGroup{}

	topicCount = flag.Int("topic_count", 1, "topic count")
	interval   = flag.Duration("interval", time.Second, "publish time interval")
)

func main() {
	flag.Parse()
	fmt.Printf("start(topic_count=%d, interval=%s)...\n", *topicCount, *interval)

	wg.Add(*topicCount)
	ns := time.Now().Nanosecond() // make sure use new topic
	for i := 0; i < *topicCount; i++ {
		topic := fmt.Sprintf("test_%d_%d", ns, i)
		if err := NewProducer(topic); err != nil {
			panic(err)
		}
	}

	// NewConsumer after topic creation
	for i := 0; i < *topicCount; i++ {
		topic := fmt.Sprintf("test_%d_%d", ns, i)
		if err := NewConsumer(topic, &TestHandler{}); err != nil {
			panic(err)
		}
	}

	wg.Wait() // wait all consumers to be ready
	fmt.Println("Init Done!!!")

	loopToPrint()
}

type TestHandler struct {
	isReady bool
}

func (th *TestHandler) HandleMessage(message *nsq.Message) error {
	// skip the first messag, which is used to create the topic
	if !th.isReady {
		th.isReady = true
		wg.Done() // one Consumer is ready
		return nil
	}

	from, err := time.Parse(time.RFC3339Nano, string(message.Body))
	if err != nil {
		return err
	}

	dur := time.Now().Sub(from)
	resultChan <- int64(dur)
	return nil
}

func loopToPrint() {
	var (
		tickerCount int64
		tickerDelay int64
		totalCount  int64
		totalDelay  int64
		maxDelay    int64
	)

	ticker := time.NewTicker(1 * time.Second)
	for {
		select {
		case delay := <-resultChan:
			tickerCount++
			tickerDelay += delay
			totalCount++
			totalDelay += delay
			if maxDelay < delay {
				maxDelay = delay
			}

		case <-ticker.C:
			if totalCount != 0 && tickerCount != 0 {
				fmt.Printf("AverageDelay: %s, \tMaxDelay: %s, \tTotalCount: %d, \tTickerCount: %d, \tTickerAverageDelay: %s\n",
					time.Duration(totalDelay/totalCount), time.Duration(maxDelay), totalCount, tickerCount, time.Duration(tickerDelay/tickerCount))
			}
			tickerCount = 0
			tickerDelay = 0
		}
	}
}

func NewConsumer(topic string, h *TestHandler) error {
	cfg := nsq.NewConfig()
	cfg.LookupdPollInterval = 10 * time.Second
	cfg.MaxInFlight = 1000
	cfg.HeartbeatInterval = 5 * time.Second
	cfg.MaxAttempts = 0

	consumer, err := nsq.NewConsumer(topic, "1", cfg)
	if err != nil {
		return err
	}

	consumer.AddConcurrentHandlers(h, 1)
	if err = consumer.ConnectToNSQLookupd("127.0.0.1:4161"); err != nil {
		return err
	}

	consumer.SetLoggerLevel(nsq.LogLevelError)

	return nil
}

func NewProducer(topic string) error {
	cfg := nsq.NewConfig()
	cfg.LookupdPollInterval = 10 * time.Second
	cfg.MaxInFlight = 1000
	cfg.HeartbeatInterval = 5 * time.Second
	cfg.MaxAttempts = 0

	p, err := nsq.NewProducer("127.0.0.1:4150", cfg)
	if err != nil {
		return err
	}
	// to create the topic
	p.Publish(topic, []byte(time.Now().Format(time.RFC3339Nano)))

	go func() {
		wg.Wait() // wait for all consumers ready
		ticker := time.NewTicker(*interval)
		for {
			select {
			case <-ticker.C:
				p.PublishAsync(topic, []byte(time.Now().Format(time.RFC3339Nano)), nil)
			}
		}
	}()

	return nil
}
